// This sketch reads IMU data from the SparkFun 5DOF IMU found at:
// http://www.sparkfun.com/commerce/product_info.php?products_id=741
// The Kalman filter implementation is directly taken from the work
// of Tom Pycke at: http://tom.pycke.be/mav/90/sparkfuns-5dof
// The AREF pin is connected to 3.3V to better use the 10 bit A/D

#include <math.h>

// Debug data 
#define BAUD 38400  

// Define pin assignments
#define LEDPIN 13
#define XACCELPIN 0
#define YACCELPIN 1
#define ZACCELPIN 2
#define XRATEPIN 4
#define YRATEPIN 3
#define ZRATEPIN 5

// Accelerometer setup
// these A/D values depend on how well the sensors are mounted
// change these values to your unique configuration
#define XMIN 405
#define XMAX 607
#define YMIN 409
#define YMAX 618
#define ZMIN 403
#define ZMAX 611
float xAccMid = ((float)(XMAX-XMIN) / 2.0) + (float)XMIN;
float yAccMid = ((float)(YMAX-YMIN) / 2.0) + (float)YMIN;
float zAccMid = ((float)(ZMAX-ZMIN) / 2.0) + (float)ZMIN;
float xAccScaleFactor = 1.0 / (float)(XMAX - XMIN) * 2.0;
float yAccScaleFactor = 1.0 / (float)(YMAX - YMIN) * 2.0;
float zAccScaleFactor = 1.0 / (float)(ZMAX - ZMIN) * 2.0;
int xAccZero, yAccZero, zAccZero;

// Gyro setup
// if AREF = 3.3V, then A/D is 931 at 3V and 465 = 1.5V
// gyro range is from +/-500 deg/sec
float GyroMid = 465;
float GyroScaleFactor = 1001 / (float)ZMAX * 0.01745328;
int xGyroZero, yGyroZero, zGyroZero;
float PiOver2 = PI / 2;

// Kalman filter
float tmp, roll_angle, pitch_angle, yaw_angle;
float dt, accX, accY, accZ;
struct Gyro1DKalman {
	float x_angle, x_bias;
	float P_00,  P_01, P_10, P_11;	
	float Q_angle, Q_gyro;
	float R_angle;
};
struct Gyro1DKalman filter_roll;
struct Gyro1DKalman filter_pitch;
struct Gyro1DKalman filter_yaw;

unsigned long previousTime = 0;
unsigned long currentTime = 0;

void setup() {
  digitalWrite(LEDPIN, HIGH);
  analogReference(EXTERNAL); // Current external ref is 3.3V
  Serial.begin(BAUD);
  init_Gyro1DKalman(&filter_roll, 0.0001, 0.0003, 0.69);
  init_Gyro1DKalman(&filter_pitch, 0.0001, 0.0003, 0.69);
  init_Gyro1DKalman(&filter_yaw, 0.0001, 0.0003, 0.69);
  previousTime = millis();
  xAccZero = xAccMid - findZero(XACCELPIN);
  yAccZero = yAccMid - findZero(YACCELPIN);
  zAccZero = ZMAX - findZero(ZACCELPIN);
  xGyroZero = GyroMid - findZero(XRATEPIN);
  yGyroZero = GyroMid - findZero(YRATEPIN);
  zGyroZero = GyroMid - findZero(ZRATEPIN);
}

void loop() {
  currentTime = millis();
  dt = (float)(currentTime - previousTime) / 1000.0;
  previousTime = currentTime;
  
  accX = accelero_x();
  accY = accelero_y();
  accZ = accelero_z();

  tmp = PredictAccG_roll(accY, accZ);
  ars_predict(&filter_roll, gyro_roll_rad(), dt);    // Kalman predict
  roll_angle = ars_update(&filter_roll, tmp);   // Kalman update + result (angle) 

  tmp = PredictAccG_pitch(accX, accZ);
  ars_predict(&filter_pitch, gyro_pitch_rad(), dt);    // Kalman predict
  pitch_angle = ars_update(&filter_pitch, tmp);   // Kalman update + result (angle) 

  tmp = PredictAccG_yaw(accX, accY);
  ars_predict(&filter_yaw, gyro_yaw_rad(), dt);    // Kalman predict
  yaw_angle = ars_update(&filter_yaw, tmp);   // Kalman update + result (angle) 

  digitalWrite(LEDPIN, HIGH);
  Serial.print((int)(accX*1000));
  Serial.print(",");
  Serial.print((int)(accY*1000));
  Serial.print(",");
  Serial.print((int)(accZ*1000));
  Serial.print(",");
  Serial.print((int)(roll_angle*1000));
  Serial.print(",");
  Serial.print((int)(pitch_angle*1000));
  Serial.print(",");
  Serial.println((int)(yaw_angle*1000));  
  digitalWrite(LEDPIN, LOW);  
}

float findZero(int channel) {
  float zero = 0;
  for (int i=0; i< 100; i++) zero += analogRead(channel);
  return (zero / 100);
}  

// roll is arctan(az/ay)-PI/2 or arctan2(az, ay)-PI/2
float PredictAccG_roll(float a_y, float a_z) {
  //return atan2(a_z, a_y)-PiOver2;
  return atan2(a_y, a_z);
}

// pitch is arctan(az/ax)-PI/2 or arctan2(az, ax)-PI/2
float PredictAccG_pitch(float a_x, float a_z) {
  //return atan2(a_z, a_x)-PiOver2;
  return atan2(-a_x, a_z);
}

// yaw is arctan(ay/ax)-PI/2 or arctan2(ay, ax)-PI/2
float PredictAccG_yaw(float a_x, float a_y) {
  //return atan2(a_x, a_y)-PiOver2;
  return atan2(a_x, a_y);
}

float accelero_x() {
  return (analogRead(XACCELPIN) - xAccMid + xAccZero) * xAccScaleFactor;
}

float accelero_y() {
  return ((analogRead(YACCELPIN) - yAccMid + yAccZero) * yAccScaleFactor);
}

float accelero_z() {
  return (analogRead(ZACCELPIN) - zAccMid + zAccZero) * zAccScaleFactor;
}

float gyro_roll_rad() {
  return ((float)analogRead(XRATEPIN) - GyroMid + xGyroZero) * GyroScaleFactor;
}

float gyro_pitch_rad() {
  return ((float)analogRead(YRATEPIN) - GyroMid + yGyroZero) * GyroScaleFactor;
}

float gyro_yaw_rad() {
  return ((float)analogRead(ZRATEPIN) - GyroMid + zGyroZero) * GyroScaleFactor;
}
  
void init_Gyro1DKalman(struct Gyro1DKalman *filterdata, float Q_angle, float Q_gyro, float R_angle) {
	filterdata->Q_angle = Q_angle;
	filterdata->Q_gyro  = Q_gyro;
	filterdata->R_angle = R_angle;
}

void ars_predict(struct Gyro1DKalman *filterdata, const float dotAngle, const float dt) {
	filterdata->x_angle += dt * (dotAngle - filterdata->x_bias);
	filterdata->P_00 +=  - dt * (filterdata->P_10 + filterdata->P_01) + filterdata->Q_angle * dt;
	filterdata->P_01 +=  - dt * filterdata->P_11;
	filterdata->P_10 +=  - dt * filterdata->P_11;
	filterdata->P_11 +=  + filterdata->Q_gyro * dt;
}

float ars_update(struct Gyro1DKalman *filterdata, const float angle_m) {
	const float y = angle_m - filterdata->x_angle;
	const float S = filterdata->P_00 + filterdata->R_angle;
	const float K_0 = filterdata->P_00 / S;
	const float K_1 = filterdata->P_10 / S;
	
	filterdata->x_angle +=  K_0 * y;
	filterdata->x_bias  +=  K_1 * y;
	filterdata->P_00 -= K_0 * filterdata->P_00;
	filterdata->P_01 -= K_0 * filterdata->P_01;
	filterdata->P_10 -= K_1 * filterdata->P_00;
	filterdata->P_11 -= K_1 * filterdata->P_01;

	return filterdata->x_angle;
}
